package com.agent.controller;

import com.agent.utils.UserContext;
import dev.langchain4j.data.document.*;
import dev.langchain4j.data.document.parser.TextDocumentParser;
import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.model.embedding.EmbeddingModel;
import dev.langchain4j.store.embedding.EmbeddingStore;
import dev.langchain4j.store.embedding.EmbeddingStoreIngestor;
import lombok.RequiredArgsConstructor;
import org.springframework.web.bind.annotation.PostMapping;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RequestParam;
import org.springframework.web.bind.annotation.RestController;
import org.springframework.web.multipart.MultipartFile;
import java.io.IOException;
import java.util.List;


/*
    管理员上传知识库
 */

@RestController
@RequiredArgsConstructor
@RequestMapping("/admin/load")
public class AdminLoadController {

    private final EmbeddingModel embeddingModel;

    private final EmbeddingStore<TextSegment> embeddingStore;

    @PostMapping(value = "/vector")
    public String load(@RequestParam("file") MultipartFile file){
        // 权限判断
        if (UserContext.getUserId()!= 1){
            return "无权限上传知识库！";
        }
        System.out.println("开始导入向量数据库...");
        try {
            // 将 MultipartFile 转换为字符串内容
            String textContent = new String(file.getBytes());
            // 创建 Document 对象
            Document document = Document.from(textContent);
            // 构建文档列表
            List<Document> documents = List.of(document);

            // 使用特定的嵌入模型和嵌入存储，将数据导入向量数据库
            EmbeddingStoreIngestor.builder()
                    .embeddingModel(embeddingModel)
                    .embeddingStore(embeddingStore)
                    .build()
                    .ingest(documents); // 将文档列表转换为向量并存储（默认使用段落分割器分割文档）

            return "文件导入成功！";
        } catch (IOException e) {
            return "文件处理失败: " + e.getMessage();
        }
    }

}
